#!/bin/python

"""
The purpose of this script is to go through all the timestamped 
directories and delete all the {rho_save, j, j_filtered} variables, 
since they're taking up a huge amount of space on the shared file 
system. I'll only keep zero-dimensional data in future
simulations; final overlap_p, final overlap_m, s, s_filtered. For the 
old  
"""

import os
import re
import pickle
from utils import is_timestamped

del_vars = ['rho_save', 'j', 'j_filtered']
trnc_vars = ['overlap_p', 'overlap_m']
    
is_trajectory = \
    lambda s: bool(re.match('trajectory_data[0-9]{1,6}.pkl',s))

def timestamped_subdirs(directory):
    """
    Isn't a lambda only because it doesn't fit on one line.
    """
    return filter(is_timestamped, os.listdir(directory))

def reduce_vars(filename, del_var_lst, trnc_var_lst):
    """
    This function takes a filename and a two lists of variables.
    We assume that the file contains a single pickled dictionary.
    This matches the storage of trajectory data.
    The first variable list indicates the keys in the dictionary 
    to be removed. 
    The second indicates variables which are to be replaced in the 
    dictionary by their last elements.
    """
    with open(filename, 'r') as phil:
        temp_dict = pickle.load(phil)
    
    for key in del_var_lst:
        try:
            del temp_dict[key]
        except KeyError:
            continue
    
    for key in trnc_var_lst:
        try:
            temp_dict[key] = temp_dict[key][-1]
        except IndexError:
            continue

    with open(filename, 'w') as phil:
        pickle.dump(temp_dict, phil)

if __name__ == '__main__':
    for sdir in timestamped_subdirs(os.getcwd()):
        os.chdir(sdir+'/pickled_output/')
        for filename in os.listdir(os.getcwd()):
            if is_trajectory(filename):
                reduce_vars(filename, del_vars, trnc_vars)
        os.chdir('../..')
    
